#if __CUDA_ARCH__ == 100
# error Atomics only used with > sm_10 architecture
#endif
#include "cuda_runtime_api.h"
#include "cuFalloc.h"
#include <malloc.h>
#include <string.h>

typedef struct __align__(8) _cuFallocHeapChunk {
	unsigned short magic;
	unsigned short count;
	volatile struct _cuFallocHeapChunk* next;
	void* reserved;
} fallocHeapChunk;

typedef struct __align__(8) _cuFallocHeap {
	void* reserved;
	size_t chunkSize;
	size_t chunks;
	size_t fallocHeapLength;
	size_t freeChunksLength;					// Size of circular buffer (set up by host)
	fallocHeapChunk** freeChunks;				// Start of circular buffer (set up by host)
	volatile fallocHeapChunk** freeChunksPtr;	// Current atomically-incremented non-wrapped offset
} fallocHeap;


///////////////////////////////////////////////////////////////////////////////
// HOST SIDE

//
//  cudaFallocInit
//
//  Takes a buffer length to allocate, creates the memory on the device and
//  returns a pointer to it for when a kernel is called. It's up to the caller
//  to free it.
//
extern "C" cudaFallocHost cudaFallocInit(size_t chunkSize, size_t length, cudaError_t* error, void* reserved) {
	cudaError_t localError; if (error == nullptr) error = &localError;
	cudaFallocHost host; memset(&host, 0, sizeof(cudaFallocHost));
	// fix up chunk to include fallocHeapChunk
	chunkSize += sizeof(fallocHeapChunk);
	if ((chunkSize % 16) > 0)
		chunkSize += 16 - (chunkSize % 16);
	// fix up length to be a multiple of chunkSize
	length = (length < chunkSize ? chunkSize : length);
	if (length % chunkSize)
		length += chunkSize - (length % chunkSize);
	size_t chunks = (size_t)(length / chunkSize);
	if (!chunks)
		return host;
	// Fix up length to include fallocHeap
	length += sizeof(fallocHeap);
	if ((length % 16) > 0)
		length += 16 - (length % 16);
	// Allocate a print buffer on the device and zero it
	fallocHeap* heap;
	if ((*error = cudaMalloc((void**)&heap, length)) != cudaSuccess)
		return host;
	cudaMemset(heap, 0, length);
	// transfer to heap
	fallocHeap hostHeap;
	hostHeap.reserved = reserved;
	hostHeap.chunkSize = chunkSize;
	hostHeap.chunks = chunks;
	hostHeap.freeChunksLength = chunks * sizeof(fallocHeapChunk*);
	hostHeap.fallocHeapLength = sizeof(fallocHeap) + freeChunksLength;
	hostHeap.freeChunks = nullptr;
	if ((*error = cudaMemcpy(heap, &hostHeap, sizeof(fallocHeap), cudaMemcpyHostToDevice)) != cudaSuccess)
		return host;
	// return the heap
	host.reserved = reserved;
	host.heap = heap;
	host.length = length;
	return host;
}

//
//  cudaFallocEnd
//
//  Frees up the memory which we allocated
//
extern "C" void cudaFallocEnd(cudaFallocHost &host) {
	if (!host.heap)
		return;
	cudaFree(host.heap); host.heap = nullptr;
}
